#include <QStringList>
#include "RestoreVisitor.h"
#include "ast/Suite.h"
#include "ast/excepthandlerType.h"
#include "CramBuilder.h"

RestoreVisitor::RestoreVisitor(CramBuilder* builder):
    VisitorBase(),
    _builder(builder),
    _inTry(0)
{

}

void RestoreVisitor::traverse(const SimpleNodePtr& node)
{
    node->traverse(this);
}

void RestoreVisitor::unhandled_node(const SimpleNodePtr& node)
{
}

void RestoreVisitor::visitSuite(const QSharedPointer<Suite>& node)
{
    node->accept(this);
}

void RestoreVisitor::visitImport(const ImportPtr& node)
{
    foreach(aliasTypePtr alias, node->names()){
        *_builder << _builder->ident << "import " << alias->name()->id();
        if (alias->asname()){
            *_builder << " as " << alias->asname()->id();
        }
        *_builder << "\n";
        //_builder->addImport(alias->name()->id(), "", alias->asname() ? alias->asname()->id() : "");
    }
}

void RestoreVisitor::visitImportFrom(const ImportFromPtr& node)
{
    if (node->names().length()){
        foreach(aliasTypePtr alias, node->names()){
            *_builder << _builder->ident << "from " << (node->module()->id().isEmpty() ? "." : node->module()->id()) << " import " << alias->name()->id();
            if ( alias->asname()){
                *_builder << " as " << alias->asname()->id();
            }
            *_builder << "\n";
            //_builder->addImport(alias->name()->id(), node->module()->id(), alias->asname() ? alias->asname()->id() : "", _inTry);
        }
    } else if (node->module()){
        *_builder << _builder->ident << "from " << node->module()->id() << " import *\n";
    }
}

void RestoreVisitor::visitTryExcept(const TryExceptPtr& node)
{
    *_builder << _builder->ident << "try:\n";
    _builder->incIdent(true);
    node->body().traverse(this);
    _builder->incIdent(false);
    foreach(excepthandlerTypePtr exp, node->handlers()){
        *_builder << _builder->ident << "except ";
        if (exp->type_())
            *_builder << makeValue(exp->type_());
        *_builder << ":\n";
        _builder->incIdent(true);
        exp->body().traverse(this);
        _builder->incIdent(false);
    }
}

void RestoreVisitor::visitClassDef(const ClassDefPtr& node)
{
    CramBuilder * builder = _builder;
    QStringList bases;
    //qDebug() << node->bases()->toString();
    foreach(exprTypePtr exp, node->bases()){
        bases.append(makeValue(exp));
    }
    foreach(keywordTypePtr exp, node->keywords()){
        if (exp->arg()->id() == "metaclass")
            bases.append("metaclass="+makeValue(exp->value()));
        //qDebug() << exp->toString();
    }
    CramBuilder* clsBuilder = builder->addClass(node->name()->id(), bases, "");

    //clsBuilder->setBases(bases);

    _builder = clsBuilder;
    node->body().traverse(this);
    _builder = builder;
}

void RestoreVisitor::visitFunctionDef(const FunctionDefPtr& node)
{
    QStringList args;
    int cnt = 0;
    foreach(exprTypePtr arg, node->args()->args()){
        QString argName = arg.dynamicCast<Name>()->id();
        if (node->args()->defaults()[cnt]){
            exprTypePtr exp = node->args()->defaults()[cnt];
            argName += "="+makeValue(exp);
        }
        args.append(argName);
        cnt++;
    }
    //fnc->setArgs(args);
    QString doc;
    foreach(stmtTypePtr after, node->body()){
        if (after->type() == "Expr" && after.dynamicCast<Expr>()->value()->type() == "Str"){
            doc += after.dynamicCast<Expr>()->value().dynamicCast<Str>()->s();
        }
    }

    _builder->addFunction(node->name()->id(), args, doc, "");
}

void RestoreVisitor::visitAssign(const AssignPtr& node)
{
    QStringList targets;
    foreach(exprTypePtr target, node->targets()){
        targets.append(makeValue(target));
    }

    *_builder << _builder->ident << targets.join(" = ") << " = " << makeValue(node->value()) << "\n";
}

QString RestoreVisitor::makeValue(const SimpleNodePtr& value)
{
    if (value->type() == "Str"){
        QString quot = "\"";
        StrPtr s = value.dynamicCast<Str>();
        if (s->type_() == StrTypeType::SingleSingle)
            quot = "\'";
        if (s->type_() == StrTypeType::SingleDouble)
            quot = "\"";
        if (s->type_() == StrTypeType::TripleDouble)
            quot = "\"\"\"";
        if (s->type_() == StrTypeType::TripleSingle)
            quot = "\'\'\'";
        return (s->binary()? "b" : "")+quot+s->s()+quot;
    }
    if (value->type() == "Num"){
        return value.dynamicCast<Num>()->num();
    }
    if (value->type() == "Name"){
        return value.dynamicCast<Name>()->id();
    }
    if (value->type() == "Tuple"){
        QStringList out;
        TuplePtr tup = value.dynamicCast<Tuple>();
        foreach(exprTypePtr exp, tup->elts()){
            out.append(makeValue(exp));
        }
        return "("+out.join(", ")+(tup->endsWithComma() ? ",": "")+")";
    }
    if (value->type() == "List"){
        QStringList out;
        ListPtr tup = value.dynamicCast<List>();
        foreach(exprTypePtr exp, tup->elts()){
            out.append(makeValue(exp));
        }
        return "("+out.join(", ")+")";
    }
    if (value->type() == "Call"){
        CallPtr call = value.dynamicCast<Call>();
        return makeValue(call->func())+"("+makeCallArgs(call->args())+")";
    }
    if (value->type() == "Attribute"){
        AttributePtr attr = value.dynamicCast<Attribute>();
        return makeValue(attr->value())+"."+attr->attr().dynamicCast<NameTok>()->id();
    }
    if (value->type() == "DictComp" || value->type() == "Dict"){
        return "{}";
    }
    if (value->type() == "UnaryOp"){
        UnaryOpPtr oper = value.dynamicCast<UnaryOp>();
        QString op = "";
        if (oper->op() == UnaryopType::Invert)
            op = "~";
        else if (oper->op() == UnaryopType::UAdd)
            op = "+";
        else if (oper->op() == UnaryopType::USub)
            op = "-";
        else if (oper->op() == UnaryopType::Not)
            op = "not ";
        return op+makeValue(oper->operand());
    }
    if (value->type() == "BinOp"){
        BinOpPtr oper = value.dynamicCast<BinOp>();
        QString op = "";
        switch(oper->op()){
        case OperatorType::Add: op = "+"; break;
        case OperatorType::Sub: op = "-"; break;
        case OperatorType::Mult: op = "*"; break;
        case OperatorType::Div: op = "/"; break;
        case OperatorType::Mod: op = "%"; break;
        case OperatorType::Pow: op = "**"; break;
        case OperatorType::LShift: op = "<<"; break;
        case OperatorType::RShift: op = ">>"; break;
        case OperatorType::BitOr: op = "|"; break;
        case OperatorType::BitXor: op = "^"; break;
        case OperatorType::BitAnd: op = "&"; break;
        case OperatorType::FloorDiv: op = "//"; break;
        }
        return makeValue(oper->left())+" "+op+" "+makeValue(oper->right());
    }
    if (value->type() == "ListComp"){
        value->toString();
        ListCompPtr val = value.dynamicCast<ListComp>();
        QStringList gen;
        foreach(comprehensionTypePtr comp, val->generators()){
            gen.append(makeValue(comp));
        }
        return "["+makeValue(val->elt())+" "+gen.join(", ")+"]";
    }
    if (value->type() == "Comprehension"){
        ComprehensionPtr comp = value.dynamicCast<Comprehension>();
        return "for "+makeValue(comp->target())+" in "+makeValue(comp->iter());
    }
    if (value->type() == "Subscript"){
        SubscriptPtr sub = value.dynamicCast<Subscript>();
        return makeValue(sub->value())+"["+makeValue(sub->slice())+"]";
    }
    if (value->type() == "Index"){
        IndexPtr ind = value.dynamicCast<Index>();
        return makeValue(ind->value());
    }
    if (value->type() == "Lambda"){
        LambdaPtr lambda = value.dynamicCast<Lambda>();
        QStringList args;
        int cnt = 0;
        foreach(exprTypePtr arg, lambda->args()->args()){
            QString argName = arg.dynamicCast<Name>()->id();
            if (lambda->args()->defaults()[cnt]){
                exprTypePtr exp = lambda->args()->defaults()[cnt];
                argName += "="+makeValue(exp);
            }
            args.append(argName);
            cnt++;
        }
        return "lambda "+args.join(", ")+": "+makeValue(lambda->body());
    }
    if (value->type() == "Compare"){
        ComparePtr compare = value.dynamicCast<Compare>();
        QStringList comps;
        foreach(exprTypePtr c, compare->comparators())
            comps.append(makeValue(c));

        QStringList ops;
        foreach(int op, compare->ops()){
            switch(op){
            case CmpopType::Eq: ops.append("=="); break;
            case CmpopType::NotEq: ops.append("!="); break;
            case CmpopType::Lt: ops.append("<"); break;
            case CmpopType::LtE: ops.append("<="); break;
            case CmpopType::Gt: ops.append(">"); break;
            case CmpopType::GtE: ops.append(">="); break;
            case CmpopType::Is: ops.append("is"); break;
            case CmpopType::IsNot: ops.append("is not"); break;
            case CmpopType::In: ops.append("in"); break;
            case CmpopType::NotIn: ops.append("not in"); break;
            }
        }
        QStringList res;
        for(int i = 0; i < ops.length(); ++i){
            res.append(ops[i]+" "+comps[i]);
        }
        return makeValue(compare->left())+" "+res.join(" ");
    }
    if (value->type() == "StrJoin"){
        StrJoinPtr sj = value.dynamicCast<StrJoin>();
        QStringList sts;
        foreach(exprTypePtr s, sj->strs()){
            sts.append(makeValue(s));
        }
        return sts.join(" + ");
    }
    qDebug() << "Unhandled value" << value->type() << value->toString();

    return "";
}

QString RestoreVisitor::makeCallArgs(const ExprTypeList& args)
{
    QStringList out;
    foreach(exprTypePtr exp, args){
        out.append(makeValue(exp));
    }
    return out.join(", ");
}

void RestoreVisitor::visitPass(const PassPtr& /*node*/)
{
    *_builder << _builder->ident << "pass\n";
}

void RestoreVisitor::visitName(const NamePtr& node)
{
    *_builder << node->id();
}

void RestoreVisitor::visitExpr(const ExprPtr& node)
{
    *_builder << _builder->ident << makeValue(node->value()) << "\n";
}

void RestoreVisitor::visitRaise(const RaisePtr& node)
{
    *_builder << _builder->ident << "raise "<< makeValue(node->type_()) << "\n";
}

void RestoreVisitor::visitIf(const IfPtr& node)
{

}
